#PBAR_ExamineMarkerValues.py
#
#Open all the markerValue files
#
#  4/25/2013, JK, added annotation
#

# plot of discriminant vs count


#PBAR_PrcoessStandardWidthImages.py

import PBAR_Zspec, PBAR_Cargo
import numpy as np
import os
import matplotlib.pyplot as plt
#from scipy import interpolate
from matplotlib import cm
from scipy import ndimage
import copy
import glob
#import json
import pickle


# Set useful plot variables
plotColors = ['r', 'b', 'g', 'm', 'c', 'y', 'k'] * 10
lineStyles = ['-', '-.', ':', '_', '|'] *  10
markerTypes = ['.', 'o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd']
markerTypes = markerTypes * 2
colorList = np.array(['r', 'b', 'g', 'm', 'c', 'y'])

# list of zpsec detector names - Z1 to Z136
ZspecNameList = np.array(['Z%d'%i for i in np.arange(1, 137)])
# list of bad detector numbers starting at 1
badZspecList = np.array([1,2,3,4,5,6,7,8,20,26,31,33,34,38,39,40,44,53,56,62,68,76,80,125,126,127,128,129,130,131,132,133,134,135,136])
# list of good detector numbers starting at 1
goodZspecList = np.array([i for i in np.arange(1, 137) if (i not in badZspecList)])
# list of bad detector names
badZspecNameList = ZspecNameList[(badZspecList-1)]
# list of good detector names
goodZspecNameList = np.array(['Z%d'%i for i in np.arange(1, 137) if (i not in badZspecList)])

badZspecIndices = badZspecList - 1
goodZspecIndices = goodZspecList - 1

badZspecMask = np.zeros(len(ZspecNameList))
badZspecMask[badZspecIndices] = True
badZspecMask = badZspecMask.astype(bool)

goodZspecMask = np.zeros(len(ZspecNameList))
goodZspecMask[goodZspecIndices] = True
goodZspecMask = goodZspecMask.astype(bool)

#dataDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3\BasicScansStandardWidth'
#plotDir = r'C:\Users\jkwong\Documents\Work\PBAR\data3\BasicScansPlots'

dataDir = r'E:\PBAR\data3b\BasicScansStandardWidth'
plotDir = r'E:\PBAR\data3b\BasicScansPlots'
basepathSet2 = r'C:\Users\jkwong\Documents\Work\PBAR\data4\Mar-files'


# define parameters for standard width images
numberTimeSlices = 2000
preCargoTimeSlices = 70

acqTime = 1/60.
figureSize = (16, 10)

params = {'legend.fontsize': 8}
plt.rcParams.update(params)

#################################
## DEFINE FILE NAMES

cargoConfigBaseDirList = glob.glob(os.path.join(dataDir, '*'))


markerValuesList = []
fullFilenameMarkerValuesList = []
fullFilenameZspecList = []
filenameZspecList = []
fullFilenameCargoList = []
filenameCargoList = []
cargoNumberList = []


for (dirIndex, cargoConfigBaseDir) in enumerate(cargoConfigBaseDirList):
    subDirList = glob.glob(os.path.join(cargoConfigBaseDir, '*'))
    for (subDirIndex, dataPath) in enumerate(subDirList):
        print(dataPath)

        a, b = os.path.split(dataPath)
        a, c = os.path.split(a)

        cargoNumberList.append(int(c))
        plotSaveDir = os.path.join(plotDir, c, b)
        
        # make paths if not present
        if not os.path.exists(os.path.join(plotDir, c)):
            os.mkdir(os.path.join(plotDir, c))

        if not os.path.exists(plotSaveDir):
            os.mkdir(plotSaveDir)
        
        # find the zspec scan number
        temp = glob.glob(os.path.join(dataPath, '*.npy'))
        # if file doesn't exist skip this set completely
        if temp == []:
            print('No npy file found.  Skipping.')
            continue
        
        for t in temp:
            (a, b) = os.path.split(t)
            if len(b) < 22:  # we need to do this now to avoid the cargoimage.npy
                a, filename = os.path.split(temp[0])
                zspecScanNumber = filename[0:4]
                filenameZspec = filename
                fullFilenameZspec = os.path.join(dataPath, filename)
                break
        
        # Get the cargo image name
        temp = glob.glob(os.path.join(dataPath, '*.cargoimage.npy'))
        # if file doesn't exist skip this set completely
        if temp == []:
            print('No cargoimage file found.  Skipping.')
            continue
        a, filename = os.path.split(temp[0])
        cargoScanNumber = filename[0:16]
        filenameCargo = filename    
        fullFilenameCargo = os.path.join(dataPath, filename)
        
        fullFilenameMarker = os.path.join(dataPath, filename.replace('cargoimage.npy', 'cargomarker'))
        fullFilenameMarkerValues = os.path.join(dataPath, filename.replace('cargoimage.npy', 'cargomarkervalues'))

        fullFilenameMarkerValuesList.append(fullFilenameMarkerValues)
        fullFilenameZspecList.append(fullFilenameZspec)
        filenameZspecList.append(filenameZspec)
        fullFilenameCargoList.append(fullFilenameCargo)
        filenameCargoList.append(filenameCargo)

        ######################
        ## READ IN DATA
        
#        # Read in marker files
#        markerStandardWidth = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
#        
#        # Read Zspec
#        datStandardWidth = np.load(fullFilenameZspec)
#        
#        # Read Cargo
#        datCargoStandardWidth= np.load(fullFilenameCargo)
        
        with open(fullFilenameMarkerValues, 'rb') as fid:
            markerValuesList.append(pickle.load(fid))


# Make nice arrays for the markerValues
#featuresBasicScan[binning type][feature name]

featuresBasicScan = {}
featuresBasicScanMean = {}

boxTypeList = ['center', 'centerBlock', 'block', 'blockMask']

featureNameList = markerValuesList[0][0]['center']['discrim'].keys()
featureNameList = ['binMean', 'binKur', 'binSkew', 'count', 'dist0']


for boxTypeIndex, boxType in enumerate(boxTypeList):
    print(boxType)
    featuresBasicScan[boxType]= {}
    featuresBasicScanMean[boxType] = {}
    
    for featureNameIndex, featureName in enumerate(featureNameList):   # cycle through features
        featuretemp = []
        featuremeantemp = []
        
        materialtemp = []
        materialmeantemp = []

        cargonumbertemp = []
        carognumbermeantemp = []
        
        # build up the feature array by concatenating all the values
        for (markerIndex, marker) in enumerate(markerValuesList):  # go through datasets
            cargoNumber = cargoNumberList[markerIndex]
            for i, mark in enumerate(marker): # go through markers in a dataset
#                print('cargo config %d, marker %s'  %(cargoNumberList[markerIndex], mark['target']))
               
                # center point
                if boxType == 'center':
                    xx = mark['center']['discrim'][featureName]
                    xxMean = mark['center']['discrim'][featureName]
                elif boxType == 'centerBlock':        
                    # center block
                    xx = mark['centerBlock']['discrim'][featureName]
                    xxMean = mark['centerBlock']['discrimMean'][featureName]
                elif boxType == 'block':  
                    # center
                    xx = mark['block']['discrim'][featureName]
                    xxMean = mark['block']['discrimMean'][featureName]
                elif boxType == 'blockMask':  
                    # center 
                    xx = mark['blockMask']['discrim'][featureName]
                    xxMean = mark['blockMask']['discrimMean'][featureName]
        
        
#                # skip if the value is nan or inf
#                if np.isnan(yy) or np.isnan(xx) or np.isinf(yy) or np.isinf(xx):
#                    print('Value is nan or inf, skipping')
#                    continue
                
                featuremeantemp.append(xxMean)
                
                try:
                    featuretemp = featuretemp + list(xx)
                except:
                    featuretemp.append(xx)
                
                if (len(mark['target']) > 0):
                    materialmeantemp.append(mark['target'][0])
                    try:
                        materialtemp = materialtemp + len(xx)*[mark['target'][0]]
                    except:
                        materialtemp = materialtemp + [mark['target'][0]]
                else:
                    materialmeantemp.append('_')
                    try:
                        materialtemp = materialtemp + len(xx)*['_']
                    except:
                        materialtemp = materialtemp + ['_']

                carognumbermeantemp.append(cargoNumber)
                try:
                    cargonumbertemp = materialtemp + len(xx)*[cargoNumber]
                except:
                    cargonumbertemp = materialtemp + [cargoNumber]

        featuresBasicScanMean[boxType][featureName] = np.array(featuremeantemp)                    
        featuresBasicScan[boxType][featureName] = np.array(featuretemp)
        
        featuresBasicScanMean[boxType]['material'] = np.array(materialmeantemp)
        featuresBasicScan[boxType]['material'] = np.array(materialtemp)

        featuresBasicScanMean[boxType]['cargoNumber'] = np.array(carognumbermeantemp)
        featuresBasicScan[boxType]['cargoNumber'] = np.array(cargonumbertemp)



#  ANALYSIS

# old values from 10/2013
removeOffset = 0
wBest = np.array([-2.07779891,  1.88095226,  4.65429156])
pfitmean = np.array([ -3.45931698e-06,   1.80445624e-03,  -3.46939502e-01,  9.03548991e-01])
pfitsigma = np.array([ -2.90915823e-07,   2.71344856e-04,  -6.81517865e-02, 6.54989380e+00])


# Plots histogram of the center and centerblock values and also the mean value
# Comparing the 'center' and 'centerblock' data points

plt.figure()
plt.grid()
feature1 = 'count'
feature1 = 'dist0'
index = 82
markerIndex = 2

hist(markerValuesList[index][markerIndex]['block']['discrim'][feature1])
plt.plot(markerValuesList[index][markerIndex]['block']['discrimMean'][feature1], 0, 'ob', markersize = 15)

hist(markerValuesList[index][markerIndex]['centerBlock']['discrim'][feature1])
plt.plot(markerValuesList[index][markerIndex]['centerBlock']['discrimMean'][feature1], 0, 'og', markersize = 15)

plt.plot(markerValuesList[index][markerIndex]['center']['discrim'][feature1], 0, 'or', markersize = 15)

plt.title(filenameZspecList[index] + ', ' + markerValuesList[index][markerIndex]['target'])
plt.legend()
plt.xlabel('feature1')
plt.ylabel('Count')

#  Comparing the dist value calculated with 'center' versus that from 'centerblock'.
#  
#  If everything is fine, there should be anticorrelation here
#  This will plot a lot of points
#

plt.figure()
plt.grid()
#feature1 = 'count'
#feature1 = 'dist0'
feature1 = 'binKur'
removeOffset = 0

markerAlpha = 0.2

for (markerIndex, marker) in enumerate(markerValuesList):  # go through datasets          
    for i, mark in enumerate(marker): # go through markers in a dataset
    
#        xx, yy = mark['center']['discrim'][feature1], mark['center']['discrim'][feature2]
        xx = copy.copy(mark['center']['discrim'][feature1])
#        yy = copy.copy(mark['centerBlock']['discrimMean'][feature1])
        yy = copy.copy(mark['blockMask']['discrimMean'][feature1])
        
        if removeOffset:
            yy = yy - np.polyval(pfitmean, xx)
        try:
            if mark['target'][0] == 'S':
                plt.plot(xx, yy, 'dy', markersize = 12, alpha  = markerAlpha, label = mark['target'] +' Center')
            elif mark['target'][0] == 'W':
                plt.plot(xx, yy, 'vb', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
            elif mark['target'][0] == 'D':
                plt.plot(xx, yy, 'sb', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
            elif mark['target'][0] == 'P':
                plt.plot(xx, yy, '*c', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
            elif mark['target'][0] == 'F':# low density stuff
                plt.plot(xx, yy, 'or', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
            else:
                plt.plot(xx, yy, 'om', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
        except:
            plt.plot(xx, yy, 'om', markersize = 12, alpha  = markerAlpha, label = mark['target']+' Center')
            
plt.title('Discriminant vs Counts, Cargos 1-48')
plt.legend()
plt.xlabel('Center, ' + feature1)
plt.ylabel('Center Block, ' + feature1)
#plt.xlim((0, 500))
#plt.ylim((-100, 100))


for (markerIndex, marker) in enumerate(markerValuesList):  # go through datasets          
    for i, mark in enumerate(marker): # go through markers in a dataset
        if mark['target'] == 'Center':
            print(markerIndex, i)



# scatter plot of the discriminant vs zspec count


plt.figure()
plt.grid()
plotAlpha = 0.3
markSize = 12
labelPoints = 1
markerAlpha = 0.2
textAlpha = 0.5

boxType = 'center'
#boxType = 'centerBlock'
#boxType = 'block'
#boxType = 'blockMask'

removeOffset = 0

feature1 = 'count'
feature2 = 'dist0'
feature2 = 'binKur'
feature2 = 'binSkew'
feature2 = 'binMean'


cargoCountRange = np.array([0.0, 0.4e8])

for (markerIndex, marker) in enumerate(markerValuesList):  # go through datasets          
    for i, mark in enumerate(marker): # go through markers in a dataset
    
        print('cargo config %d, marker %s'  %(cargoNumberList[markerIndex], mark['target']))
       
#       if cargoNumberList[markerIndex] < 13:
#            print('Unsure gain calibration. Skipping.')
#            continue
    
        # center point
        if boxType == 'center':
            xx = mark['center']['discrim'][feature1]
            yy = mark['center']['discrim'][feature2]
            inRange = ( mark['center']['discrim']['cargoCount'] >= cargoCountRange[0] ) & \
                ( mark['center']['discrim']['cargoCount'] <= cargoCountRange[1] )

        elif boxType == 'centerBlock':        
            # center block
            xx = mark['centerBlock']['discrimMean'][feature1]
            yy = mark['centerBlock']['discrimMean'][feature2]        
            inRange = ( mark['centerBlock']['discrimMean']['cargoCount'] >= cargoCountRange[0] ) & \
                ( mark['centerBlock']['discrimMean']['cargoCount'] <= cargoCountRange[1] )
        elif boxType == 'block':  
            # center 
            xx = mark['block']['discrimMean'][feature1]
            yy = mark['block']['discrimMean'][feature2]        
            inRange = ( mark['block']['discrimMean']['cargoCount'] >= cargoCountRange[0] ) & \
                ( mark['block']['discrimMean']['cargoCount'] <= cargoCountRange[1] )
        elif boxType == 'blockMask':  
            # center 
            xx = mark['blockMask']['discrimMean'][feature1]
            yy = mark['blockMask']['discrimMean'][feature2]        
            inRange = ( mark['blockMask']['discrimMean']['cargoCount'] >= cargoCountRange[0] ) & \
                ( mark['blockMask']['discrimMean']['cargoCount'] <= cargoCountRange[1] )            

        if ~inRange:
            continue
        
        if removeOffset:
            yy = yy - np.polyval(pfitmean, xx)
        
        # skip if the value is nan or inf
        if np.isnan(yy) or np.isnan(xx) or np.isinf(yy) or np.isinf(xx):
            print('Value is nan or inf, skipping')
            continue
        
        try:
            if mark['target'][0] == 'S':
                markerType = 'dy'
            elif mark['target'][0] == 'W':
                markerType = 'vb'
            elif mark['target'][0] == 'D':
                markerType = 'sb'
            elif mark['target'][0] == 'P':
                markerType = '*c'
            elif mark['target'][0] == 'F':# low density stuff
                markerType = 'or'
            else:
                markerType = 'om'
        except:
            markerType = 'om'

        plt.plot(xx, yy, markerType, markersize = markSize, alpha = plotAlpha, label = mark['target']+ ', '+ boxType)
        
        if cargoNumberList[markerIndex] == 29:
            if mark['target'][0] == 'F':
                plt.plot(xx, yy, markerType, markersize = 18, alpha = markerAlpha, label = mark['target']+ ', '+ boxType)
                if ~np.isnan(yy) and ~np.isnan(xx):
                    plt.text(xx, yy, mark['target'], color = 'g', fontsize = 16)
                print(xx, yy)
        # print number next to data point
        if labelPoints:
            plt.text(xx, yy, cargoNumberList[markerIndex], fontsize = 10, color = markerType[-1], alpha = textAlpha)
        
plt.title('Discriminant vs Counts, Cargos 1-48, %s, Remove offset %d' %(boxType, removeOffset))
plt.legend()
plt.xlabel(feature1)
plt.ylabel(feature2)
plt.xlim((0, 500))

if feature2 == 'dist0':
    plt.ylim((-25, 75))
elif feature2 == 'binKur':
    plt.ylim((0, 15))
elif feature2 == 'binSkew':
    plt.ylim((0, 4))





# zspec count vs cargo count

# scatter plot of the discriminant vs zspec count

plt.figure()
plt.grid()
feature1 = 'cargoCount'
feature2 = 'count'

for (markerIndex, marker) in enumerate(markerValuesList):  # go through datasets          
    for i, mark in enumerate(marker): # go through markers in a dataset
    
        xx = mark['center']['discrim'][feature1]
        yy = mark['center']['discrim'][feature2]
#        xx, yy = mark['centerBlock']['discrimMean'][feature1], mark['centerBlock']['discrimMean'][feature2]

        if removeOffset:
            yy = yy - np.polyval(pfitmean, xx)
        try:
            if mark['target'][0] == 'S':
                plt.plot(xx, yy, 'dy', markersize = 12, alpha  = 1.0, label = mark['target'] +' Center')
            elif mark['target'][0] == 'W':
                plt.plot(xx, yy, 'vb', markersize = 12, alpha  = 1.0, label = mark['target']+' Center')
            elif mark['target'][0] == 'D':
                plt.plot(xx, yy, 'sb', markersize = 12, alpha  = 1.0, label = mark['target']+' Center')
            elif mark['target'][0] == 'P':
                plt.plot(xx, yy, '*c', markersize = 12, alpha  = 1.0, label = mark['target']+' Center')
            elif mark['target'][0] == 'F':# low density stuff
                plt.plot(xx, yy, 'or', markersize = 12, alpha  = 1.0, label = mark['target']+' Center')
            else:
                plt.plot(xx, yy, 'om', markersize = 12, alpha  = 1.0, label = mark['target']+' Center')
        except:
            plt.plot(xx, yy, 'om', markersize = 12, alpha  = 1.0, label = mark['target']+' Center')
            
plt.title('zspec counts vs cargo counts, Cargos 1-39')
plt.legend()
plt.xlabel(feature1)
plt.ylabel(feature2)
#plt.xlim((0, 500))
#plt.ylim((-100, 100))
